from sqlalchemy.orm import Session
from models import TreeData
from database import engine
import pandas as pd
import numpy as np

def init_db():
    from database import Base
    # 删除现有表并重新创建
    Base.metadata.drop_all(bind=engine)
    Base.metadata.create_all(bind=engine)
    print("数据库表结构已重新创建")
    
    db = Session(autocommit=False, autoflush=False, bind=engine)
    
    # 读取Excel文件
    print("开始读取Excel文件...")
    df = pd.read_excel('trees.xlsx')
    print(f"Excel文件读取完成，共 {len(df)} 行数据")
    
    # 获取标题行（第1行）和数据行（从第2行开始）
    headers = df.iloc[0].tolist()
    data_rows = df.iloc[1:].reset_index(drop=True)
    print(f"标题行: {headers}")
    print(f"数据行数: {len(data_rows)}")
    
    # 创建列名映射
    column_mapping = {
        '样木号': 'sample_id',
        '检尺类型': 'inspection_type', 
        '树种': 'species',
        '2024年胸径': 'diameter_2024',
        '2025年\n胸径': 'diameter_2025',
        '方位角': 'azimuth',
        '水平距离': 'horizontal_distance',
        '大样地方位角': 'plot_azimuth',
        '大样地\n水平距离': 'plot_horizontal_distance',
        'X': 'x',
        'Y': 'y',
        '备注': 'note'
    }
    
    # 重命名列
    data_rows.columns = headers
    print(f"列名重命名完成: {list(data_rows.columns)}")
    
    # 处理数据并插入数据库
    success_count = 0
    error_count = 0
    skipped_count = 0
    duplicate_count = 0
    processed_sample_ids = set()
    
    for index, row in data_rows.iterrows():
        # 跳过空行（没有样木号的行）
        if pd.isna(row['样木号']):
            print(f"跳过空行: 第 {index+2} 行")
            skipped_count += 1
            continue
        
        sample_id = int(row['样木号'])
        
        # 检查是否已处理过该样木号
        if sample_id in processed_sample_ids:
            print(f"跳过重复样木号: {sample_id} (第 {index+2} 行)")
            duplicate_count += 1
            continue
            
        tree_data = {}
        
        # 映射数据字段
        for excel_col, db_col in column_mapping.items():
            value = row[excel_col]
            
            # 处理数值类型转换
            if db_col in ['sample_id']:
                tree_data[db_col] = int(value) if pd.notna(value) else None
            elif db_col in ['diameter_2024', 'diameter_2025', 'azimuth', 'horizontal_distance', 
                           'plot_azimuth', 'plot_horizontal_distance', 'x', 'y']:
                tree_data[db_col] = float(value) if pd.notna(value) else None
            else:
                tree_data[db_col] = str(value) if pd.notna(value) else None
        
        # 创建TreeData对象并添加到数据库
        try:
            tree_record = TreeData(**tree_data)
            db.add(tree_record)
            processed_sample_ids.add(sample_id)
            success_count += 1
            if success_count % 10 == 0:  # 每10条记录打印一次进度
                print(f"已处理 {success_count} 条记录...")
        except Exception as e:
            error_count += 1
            print(f"添加记录失败 (样木号 {tree_data.get('sample_id', 'unknown')}): {e}")
            continue
    
    try:
        db.commit()
        print(f"成功导入 {success_count} 条记录到数据库")
        print(f"失败 {error_count} 条记录")
        print(f"跳过 {skipped_count} 条记录（空行）")
        print(f"跳过 {duplicate_count} 条记录（重复样木号）")
        print(f"数据库总记录数: {db.query(TreeData).count()}")
    except Exception as e:
        db.rollback()
        print(f"提交失败: {e}")
    
    db.close()

if __name__ == "__main__":
    init_db()